package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os/exec"
	"os/user"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
)

// CPUInfo 表示CPU信息
type CPUInfo struct {
	CPU   int
	Core  int
	NUMA  int
}

// OSDInfo 表示OSD信息
type OSDInfo struct {
	ID   int
	Path string
}

// SSHConfig 表示SSH配置
type SSHConfig struct {
	User    string
	Host    string
	Port    int
	KeyPath string
}

func main() {
	var (
		numaList    = flag.String("numa", "", "NUMA节点列表，以逗号分隔，例如: 0,1,2")
		osdCount    = flag.Int("osd-count", 0, "OSD数量（默认自动发现所有OSD）")
		osdIDs      = flag.String("osd-ids", "", "指定OSD ID列表，以逗号分隔，例如: 12,13,14")
		osdCPUMap   = flag.String("osd-cpu-map", "", "指定OSD到CPU的映射，格式: OSD_ID:CPU_LIST,OSD_ID:CPU_LIST 例如: 12:1-5,13:10-15")
		fsID        = flag.String("fsid", "", "Ceph FSID (留空自动检测，使用 'cephadm shell ceph -s')")
		coresPerOSD = flag.Int("cores-per-osd", 5, "每个OSD分配的核心数")
		apply       = flag.Bool("apply", false, "是否应用更改")
		restart     = flag.Bool("restart", false, "是否重启OSD服务")
		restartMethod = flag.String("restart-method", "systemd", "重启方式: systemd|cephadm|orch (优先顺序或单一方式)")
		sudoFlag      = flag.Bool("sudo", false, "远程或本地执行管理命令时使用 sudo")
		excludeCPUs = flag.String("exclude-cpus", "96-191", "要排除的CPU列表，支持范围格式如96-191")
		remoteHost  = flag.String("H", "", "远程主机地址，使用SSH连接")
		sshKeyPath  = flag.String("ssh-key", "", "SSH私钥路径")
		sshUser     = flag.String("ssh-user", "", "SSH用户名")
		sshPort     = flag.Int("ssh-port", 22, "SSH端口")
	)
	
	flag.Parse()
	
	if *numaList == "" && *osdCPUMap == "" {
		log.Fatal("必须指定-numa参数或-osd-cpu-map参数")
	}
	
	// SSH配置
	var sshConfig *SSHConfig
	if *remoteHost != "" {
		// 设置默认SSH用户
		sshUser := *sshUser
		if sshUser == "" {
			currentUser, err := user.Current()
			if err == nil {
				sshUser = currentUser.Username
			} else {
				sshUser = "root"
			}
		}
		
		// 设置默认SSH密钥路径
		keyPath := *sshKeyPath
		if keyPath == "" {
			currentUser, _ := user.Current()
			if currentUser != nil {
				keyPath = filepath.Join(currentUser.HomeDir, ".ssh", "id_rsa")
			}
		}
		
		sshConfig = &SSHConfig{User: sshUser, Host: *remoteHost, Port: *sshPort, KeyPath: keyPath}
	}
	
	// 解析NUMA节点列表
	numas := parseNumaList(*numaList)
	if *numaList != "" {
		fmt.Printf("使用NUMA节点: %v\n", numas)
	}
	
	var allocation map[OSDInfo][]int
	var numaAssignment map[OSDInfo]int
	
	// 自动探测 FSID（如果未提供且未使用手动映射）
	if *fsID == "" && *osdCPUMap == "" {
		fsDetected, err := detectFSID(sshConfig)
		if err != nil {
			log.Printf("警告: 自动检测FSID失败: %v", err)
			log.Printf("请通过 --fsid 显式指定，例如: --fsid=<cluster-id>")
		} else {
			*fsID = fsDetected
			log.Printf("自动检测到 FSID: %s", *fsID)
		}
	}

	// 如果指定了OSD到CPU的映射，则使用该映射
	if *osdCPUMap != "" {
		if *fsID == "" { // 仍空则尝试探测
			if fsDetected, err := detectFSID(sshConfig); err == nil {
				*fsID = fsDetected
				log.Printf("自动检测到 FSID: %s", *fsID)
			}
		}
		allocation, numaAssignment = parseOSDCPUMap(*osdCPUMap, *fsID)
		fmt.Println("使用指定的OSD-CPU映射:")
		printAllocation(allocation, numaAssignment)
	} else {
		// 获取CPU拓扑信息
		cpuInfos, err := getCPUInfo(sshConfig)
		if err != nil {
			log.Fatalf("获取CPU信息失败: %v", err)
		}
		
		// 排除指定的CPU
		excluded := parseCPUList(*excludeCPUs)
		cpuInfos = filterCPUs(cpuInfos, excluded)
		
		// 按NUMA节点分组CPU
		numaCPUs := groupCPUsByNUMA(cpuInfos, numas)
		
		// 获取OSD信息
		osds, err := getOSDInfo(*fsID, *osdCount, *osdIDs, sshConfig)
		if err != nil {
			log.Fatalf("获取OSD信息失败: %v", err)
		}
		if len(osds) == 0 {
			log.Printf("未发现任何 OSD 目录。请确认 FSID (%s) 是否正确或使用 --fsid 指定。", *fsID)
			possible := listPossibleFSIDs(sshConfig)
			if len(possible) > 0 {
				log.Printf("可能的 FSID 候选: %s", strings.Join(possible, ","))
			}
		}
		
		// 为每个OSD分配CPU，确保每个OSD只使用一个NUMA节点
		allocation, numaAssignment = allocateCPUsPerNUMA(numaCPUs, osds, *coresPerOSD)
		
		// 打印分配结果
		fmt.Println("\nCPU分配结果:")
		printAllocation(allocation, numaAssignment)
	}
	
	// 如果指定了-apply，则应用更改
	if *apply {
		applyChanges(allocation, numaAssignment, sshConfig)
	}
	
	// 如果指定了-restart，则重启OSD服务
	if *restart {
		// 从allocation中提取OSD信息
		var osds []OSDInfo
		for osd := range allocation {
			osds = append(osds, osd)
		}
		restartOSDServices(osds, sshConfig, *fsID, *restartMethod, *sudoFlag)
	}
}

// detectFSID 通过 cephadm shell ceph -s 自动检测 FSID
func detectFSID(sshConfig *SSHConfig) (string, error) {
	cmds := []string{
		"cephadm shell ceph -s",
		"cephadm shell -- ceph -s",
		"sudo cephadm shell ceph -s",
		"sudo cephadm shell -- ceph -s",
		"ceph -s", // 直接尝试
		"sudo ceph -s",
	}
	for _, c := range cmds {
		out, err := executeCommand(sshConfig, c)
		if err != nil || len(out) == 0 {
			continue
		}
		lines := strings.Split(string(out), "\n")
		for _, line := range lines {
			line = strings.TrimSpace(line)
			// 典型格式: id: 58b1... 或 cluster:
			if strings.HasPrefix(line, "id:") {
				parts := strings.SplitN(line, "id:", 2)
				fs := strings.TrimSpace(parts[1])
				if fs != "" {
					return fs, nil
				}
			}
		}
	}
	return "", fmt.Errorf("无法通过 ceph -s 解析 FSID")
}

// listPossibleFSIDs 简单列出 /var/lib/ceph 下可能的 FSID 目录名称
func listPossibleFSIDs(sshConfig *SSHConfig) []string {
	out, err := executeCommand(sshConfig, "ls -1 /var/lib/ceph 2>/dev/null || echo")
	if err != nil {
		return nil
	}
	var res []string
	for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
		line = strings.TrimSpace(line)
		if line == "" { continue }
		// 粗略判断: FSID 通常是 UUID 形式，包含多处 '-'，长度 > 30
		if strings.Count(line, "-") >= 4 && len(line) > 30 {
			res = append(res, line)
		}
	}
	return res
}

// parseNumaList 解析NUMA节点列表
func parseNumaList(numaList string) []int {
	var numas []int
	parts := strings.Split(numaList, ",")
	for _, part := range parts {
		if num, err := strconv.Atoi(strings.TrimSpace(part)); err == nil {
			numas = append(numas, num)
		}
	}
	return numas
}

// parseCPUList 解析CPU列表，支持范围格式如"96-191"
func parseCPUList(cpuList string) map[int]bool {
	excluded := make(map[int]bool)
	parts := strings.Split(cpuList, ",")
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if strings.Contains(part, "-") {
			// 范围格式
			rangeParts := strings.Split(part, "-")
			if len(rangeParts) == 2 {
				start, _ := strconv.Atoi(rangeParts[0])
				end, _ := strconv.Atoi(rangeParts[1])
				for i := start; i <= end; i++ {
					excluded[i] = true
				}
			}
		} else {
			// 单个CPU
			if num, err := strconv.Atoi(part); err == nil {
				excluded[num] = true
			}
		}
	}
	return excluded
}

// buildSSHBase 构造 ssh 基础命令 (不含要执行的远程指令)
func buildSSHBase(cfg *SSHConfig) []string {
	args := []string{}
	if cfg.Port != 0 {
		args = append(args, "-p", fmt.Sprintf("%d", cfg.Port))
	}
	if cfg.KeyPath != "" {
		args = append(args, "-i", cfg.KeyPath)
	}
	// 禁用 host key 检查以避免交互（可按需开启）
	args = append(args, "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null")
	args = append(args, fmt.Sprintf("%s@%s", cfg.User, cfg.Host))
	return args
}

// executeCommand 在本地或远程执行命令
func executeCommand(sshConfig *SSHConfig, command string) ([]byte, error) {
	if sshConfig != nil && sshConfig.Host != "" {
		base := buildSSHBase(sshConfig)
		full := append([]string{}, base...)
		full = append(full, command)
		cmd := exec.Command("ssh", full...)
		return cmd.CombinedOutput()
	}
	cmd := exec.Command("/bin/sh", "-c", command)
	return cmd.CombinedOutput()
}

// readFile 读取本地或远程文件
func readFile(sshConfig *SSHConfig, filePath string) ([]byte, error) {
	if sshConfig != nil && sshConfig.Host != "" {
		cmdStr := fmt.Sprintf("cat %s", filePath)
		out, err := executeCommand(sshConfig, cmdStr)
		if err != nil {
			return nil, fmt.Errorf("远程读取文件失败: %v (%s)", err, string(out))
		}
		return out, nil
	}
	return ioutil.ReadFile(filePath)
}

// writeFile 写入本地或远程文件
func writeFile(sshConfig *SSHConfig, filePath string, data []byte) error {
	if sshConfig != nil && sshConfig.Host != "" {
		dir := filepath.Dir(filePath)
		// 创建目录
		if _, err := executeCommand(sshConfig, fmt.Sprintf("mkdir -p %s", dir)); err != nil {
			log.Printf("警告: 远程创建目录失败: %v", err)
		}
		// 通过 heredoc 写入
		escaped := strings.ReplaceAll(string(data), "EOF", "EOF_MARK") // 简单避免冲突
		cmdStr := fmt.Sprintf("cat > %s <<'EOF'\n%s\nEOF", filePath, escaped)
		if out, err := executeCommand(sshConfig, cmdStr); err != nil {
			return fmt.Errorf("远程写入文件失败: %v (%s)", err, string(out))
		}
		return nil
	}
	return ioutil.WriteFile(filePath, data, 0644)
}

// getCPUInfo 获取CPU拓扑信息
func getCPUInfo(sshConfig *SSHConfig) ([]CPUInfo, error) {
	// 执行lscpu命令获取CPU信息
	output, err := executeCommand(sshConfig, "lscpu -p=CPU,CORE,NODE")
	if err != nil {
		return nil, err
	}
	
	var cpuInfos []CPUInfo
	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		
		parts := strings.Split(line, ",")
		if len(parts) >= 3 {
			cpu, _ := strconv.Atoi(parts[0])
			core, _ := strconv.Atoi(parts[1])
			numa, _ := strconv.Atoi(parts[2])
			cpuInfos = append(cpuInfos, CPUInfo{CPU: cpu, Core: core, NUMA: numa})
		}
	}
	
	return cpuInfos, nil
}

// filterCPUs 过滤掉排除的CPU
func filterCPUs(cpuInfos []CPUInfo, excluded map[int]bool) []CPUInfo {
	var filtered []CPUInfo
	for _, cpu := range cpuInfos {
		if !excluded[cpu.CPU] {
			filtered = append(filtered, cpu)
		}
	}
	return filtered
}

// groupCPUsByNUMA 按NUMA节点分组CPU
func groupCPUsByNUMA(cpuInfos []CPUInfo, numas []int) map[int][]CPUInfo {
	// 创建NUMA节点映射
	numaMap := make(map[int]bool)
	for _, n := range numas {
		numaMap[n] = true
	}
	
	// 按NUMA节点分组
	numaCPUs := make(map[int][]CPUInfo)
	for _, cpu := range cpuInfos {
		// 只包含指定的NUMA节点
		if numaMap[cpu.NUMA] {
			numaCPUs[cpu.NUMA] = append(numaCPUs[cpu.NUMA], cpu)
		}
	}
	
	return numaCPUs
}

// parseOSDIDs 解析OSD ID列表
func parseOSDIDs(osdList string) []int {
	var osds []int
	parts := strings.Split(osdList, ",")
	for _, part := range parts {
		if num, err := strconv.Atoi(strings.TrimSpace(part)); err == nil {
			osds = append(osds, num)
		}
	}
	return osds
}

// getOSDInfo 获取OSD信息
func getOSDInfo(fsID string, count int, osdIDs string, sshConfig *SSHConfig) ([]OSDInfo, error) {
	basePath := "/var/lib/ceph/" + fsID
	
	// 如果指定了OSD ID列表，则使用指定的列表
	if osdIDs != "" {
		var osds []OSDInfo
		osdList := parseOSDIDs(osdIDs)
		for _, id := range osdList {
			osd := OSDInfo{
				ID:   id,
				Path: fmt.Sprintf("%s/osd.%d", basePath, id),
			}
			osds = append(osds, osd)
		}
		return osds, nil
	}
	
	// 自动发现OSD目录
	output, err := executeCommand(sshConfig, "ls -d "+basePath+"/osd.* 2>/dev/null || echo")
	if err != nil {
		return nil, fmt.Errorf("无法查找OSD目录: %v", err)
	}
	
	osdDirs := strings.Split(strings.TrimSpace(string(output)), "\n")
	
	// 解析OSD ID
	var osdIDsList []int
	for _, dir := range osdDirs {
		// 提取OSD ID
		parts := strings.Split(strings.TrimSpace(dir), "/")
		if len(parts) > 0 {
			osdPart := parts[len(parts)-1] // osd.X
			if strings.HasPrefix(osdPart, "osd.") {
				idStr := strings.TrimPrefix(osdPart, "osd.")
				if id, err := strconv.Atoi(idStr); err == nil {
					osdIDsList = append(osdIDsList, id)
				}
			}
		}
	}
	
	// 排序OSD ID
	sort.Ints(osdIDsList)
	
	// 如果指定了OSD数量，则只取前N个
	if count > 0 && count < len(osdIDsList) {
		osdIDsList = osdIDsList[:count]
	}
	
	// 构造OSD信息
	var osds []OSDInfo
	for _, id := range osdIDsList {
		osd := OSDInfo{
			ID:   id,
			Path: fmt.Sprintf("%s/osd.%d", basePath, id),
		}
		osds = append(osds, osd)
	}
	
	return osds, nil
}

// allocateCPUsPerNUMA 为每个OSD分配CPU，确保每个OSD只使用一个NUMA节点
func allocateCPUsPerNUMA(numaCPUs map[int][]CPUInfo, osds []OSDInfo, coresPerOSD int) (map[OSDInfo][]int, map[OSDInfo]int) {
	allocation := make(map[OSDInfo][]int)
	numaAssignment := make(map[OSDInfo]int)
	
	// 为每个NUMA节点创建核心映射
	numaCoreMap := make(map[int]map[int][]int) // numa -> core -> cpus
	for numa, cpus := range numaCPUs {
		numaCoreMap[numa] = make(map[int][]int)
		for _, cpu := range cpus {
			numaCoreMap[numa][cpu.Core] = append(numaCoreMap[numa][cpu.Core], cpu.CPU)
		}
	}
	
	// 对每个NUMA节点的核心进行排序，并过滤掉不包含有效CPU的核心
	numaSortedCores := make(map[int][]int) // numa -> sorted cores
	for numa, coreMap := range numaCoreMap {
		var cores []int
		for core, cpus := range coreMap {
			// 过滤掉不包含有效CPU的核心（即核心中的所有CPU都被排除了）
			validCPUs := 0
			for _, cpu := range cpus {
				if cpu != 0 {
					validCPUs++
				}
			}
			if validCPUs > 0 {
				cores = append(cores, core)
			}
		}
		sort.Ints(cores)
		numaSortedCores[numa] = cores
	}
	
	// 按NUMA节点循环分配OSD
	numaList := make([]int, 0, len(numaCPUs))
	for numa := range numaCPUs {
		numaList = append(numaList, numa)
	}
	sort.Ints(numaList)
	
	// 初始化每个NUMA节点的核心索引
	coreIndexPerNuma := make(map[int]int) // 每个NUMA节点的核心索引
	
	// 为每个OSD分配资源
	for i, osd := range osds {
		// 循环使用NUMA节点
		currentNumaIndex := i % len(numaList)
		currentNuma := numaList[currentNumaIndex]
		
		// 为该OSD分配核心
		var osdCPUs []int
		
		// 获取当前NUMA节点的核心列表
		sortedCores := numaSortedCores[currentNuma]
		
		// 获取当前NUMA节点的核心索引
		coreIndex := coreIndexPerNuma[currentNuma]
		
		// 分配指定数量的核心
		coresAssigned := 0
		originalCoreIndex := coreIndex
		// 第一轮分配：从当前索引开始
		for coresAssigned < coresPerOSD && coreIndex < len(sortedCores) {
			coreID := sortedCores[coreIndex]
			cpus := numaCoreMap[currentNuma][coreID]
			
			// 添加核心中的所有CPU（包括超线程），但跳过0号CPU
			for _, cpu := range cpus {
				if cpu != 0 {
					osdCPUs = append(osdCPUs, cpu)
				}
			}
			
			coresAssigned++
			coreIndex++
		}
		
		// 如果没有分配到足够的核心，从头开始继续分配（避免重复）
		if coresAssigned < coresPerOSD {
			// 从头开始分配剩余所需的核心，但要避免重复添加已有的CPU
			remainingCores := coresPerOSD - coresAssigned
			coreIndex = 0
			// 第二轮分配：从头开始直到originalCoreIndex
			for remainingCores > 0 && coreIndex < originalCoreIndex {
				coreID := sortedCores[coreIndex]
				cpus := numaCoreMap[currentNuma][coreID]
				
				// 添加核心中的所有CPU（包括超线程），但跳过0号CPU和已添加的CPU
				addedCPUs := 0
				for _, cpu := range cpus {
					// 检查是否已经添加过这个CPU
					found := false
					for _, existingCPU := range osdCPUs {
						if existingCPU == cpu {
							found = true
							break
						}
					}
					if !found && cpu != 0 {
						osdCPUs = append(osdCPUs, cpu)
						addedCPUs++
					}
				}
				
				// 只有当我们实际添加了CPU时才减少剩余核心数
				if addedCPUs > 0 {
					remainingCores--
				}
				coreIndex++
			}
		}
		
		// 更新核心索引
		coreIndexPerNuma[currentNuma] = coreIndex
		
		// 如果仍然没有分配到CPU，则至少分配一个
		if len(osdCPUs) == 0 && len(sortedCores) > 0 {
			coreID := sortedCores[0]
			cpus := numaCoreMap[currentNuma][coreID]
			for _, cpu := range cpus {
				if cpu != 0 {
					osdCPUs = append(osdCPUs, cpu)
					break
				}
			}
			// 如果除了0号CPU外没有其他CPU，那就用0号CPU
			if len(osdCPUs) == 0 && len(cpus) > 0 {
				osdCPUs = append(osdCPUs, cpus[0])
			}
		}
		
		allocation[osd] = osdCPUs
		numaAssignment[osd] = currentNuma
	}
	
	return allocation, numaAssignment
}

// printAllocation 打印CPU分配结果
func printAllocation(allocation map[OSDInfo][]int, numaAssignment map[OSDInfo]int) {
	fmt.Println("\nCPU分配结果:")
	fmt.Printf("%-6s %-10s %-20s %s\n", "OSD", "NUMA", "CPUs", "说明")
	
	for osd, cpus := range allocation {
		if len(cpus) > 0 {
			cpuStr := formatCPUList(cpus)
			numaNode := numaAssignment[osd]
			fmt.Printf("OSD %-3d %-10d %-20s %s\n", osd.ID, numaNode, cpuStr, "同一NUMA节点")
		}
	}
}


// formatCPUList 格式化CPU列表
func formatCPUList(cpus []int) string {
	if len(cpus) == 0 {
		return ""
	}
	
	// 排序
	sort.Ints(cpus)
	
	// 简单格式化（不实现完整的范围压缩）
	var parts []string
	for _, cpu := range cpus {
		parts = append(parts, strconv.Itoa(cpu))
	}
	
	// 简化处理，只显示前几个和后几个
	if len(parts) > 10 {
		return strings.Join(parts[:5], ",") + ",...," + strings.Join(parts[len(parts)-5:], ",")
	}
	
	return strings.Join(parts, ",")
}

// formatNUMAList 格式化NUMA节点列表
func formatNUMAList(numas []int) string {
	var parts []string
	for _, numa := range numas {
		parts = append(parts, strconv.Itoa(numa))
	}
	return strings.Join(parts, ",")
}

// applyChanges 应用CPU绑定更改
func applyChanges(allocation map[OSDInfo][]int, numaAssignment map[OSDInfo]int, sshConfig *SSHConfig) {
	fmt.Println("\n应用更改...")
	
	for osd, cpus := range allocation {
		unitRunPath := osd.Path + "/unit.run"
		
		// 读取文件内容
		content, err := readFile(sshConfig, unitRunPath)
		if err != nil {
			fmt.Printf("错误: 无法读取 OSD %d 的 unit.run 文件: %v\n", osd.ID, err)
			continue
		}
		
		lines := strings.Split(string(content), "\n")
		modified := false
		
		// 查找并修改包含podman run的行
		for i, line := range lines {
			if strings.Contains(line, "podman run") && strings.Contains(line, "/usr/bin/ceph-osd") {
				// 获取NUMA节点信息（通过参数传递而不是重新计算）
				// 获取分配给该OSD的NUMA节点
				numaNode := numaAssignment[osd]
				cpuStr := formatCPUList(cpus)
				
				// 移除现有的--cpuset-cpus和--cpuset-mems参数
				line = removeExistingCPUSetParams(line)
				
				// 添加新的参数，确保CPU和内存都在同一NUMA节点
				insert := fmt.Sprintf(" --cpuset-cpus=%s --cpuset-mems=%d", cpuStr, numaNode)
				if strings.Contains(line, "quay.io") {
					// 在镜像名称后插入
					line = insertAfterImageName(line, insert)
				} else {
					// 在--name参数后插入
					line = insertAfterNameParam(line, insert)
				}
				
				lines[i] = line
				modified = true
				break
			}
		}
		
		if modified {
			// 创建备份
			backupPath := unitRunPath + ".backup"
			
			// 创建备份
			err = writeFile(sshConfig, backupPath, content)
			if err != nil {
				fmt.Printf("警告: 无法创建 OSD %d 的备份文件: %v\n", osd.ID, err)
			}
			
			// 写入修改后的内容
			newContent := strings.Join(lines, "\n")
			err = writeFile(sshConfig, unitRunPath, []byte(newContent))
			if err != nil {
				fmt.Printf("错误: 无法写入 OSD %d 的 unit.run 文件: %v\n", osd.ID, err)
			} else {
				numaNode := numaAssignment[osd]
				cpuStr := formatCPUList(cpus)
				if sshConfig != nil {
					fmt.Printf("已更新 OSD %d: CPUs=%s, NUMA=%d (远程主机: %s)\n", osd.ID, cpuStr, numaNode, sshConfig.Host)
				} else {
					fmt.Printf("已更新 OSD %d: CPUs=%s, NUMA=%d\n", osd.ID, cpuStr, numaNode)
				}
			}
		} else {
			fmt.Printf("警告: 未在 OSD %d 的 unit.run 文件中找到可修改的行\n", osd.ID)
		}
	}
	
	fmt.Println("更改已应用。请手动重启OSD服务以使更改生效。")
}

// removeExistingCPUSetParams 移除现有的cpuset参数
func removeExistingCPUSetParams(line string) string {
	// 移除--cpuset-cpus参数
	line = removeParam(line, "--cpuset-cpus")
	// 移除--cpuset-mems参数
	line = removeParam(line, "--cpuset-mems")
	return line
}

// removeParam 移除指定参数
func removeParam(line, param string) string {
	for {
		start := strings.Index(line, param)
		if start == -1 {
			break
		}
		
		// 找到参数值的结束位置
		end := start + len(param)
		if end < len(line) && line[end] == '=' {
			end++ // 跳过=
			// 找到值的结束（空格或行尾）
			for end < len(line) && line[end] != ' ' {
				end++
			}
		}
		
		// 移除参数
		line = line[:start] + line[end:]
	}
	
	// 清理多余的空格
	line = strings.Join(strings.Fields(line), " ")
	return line
}

// insertAfterImageName 在镜像名称后插入内容
func insertAfterImageName(line, insert string) string {
	// 查找quay.io后的内容
	start := strings.Index(line, "quay.io")
	if start == -1 {
		return line + insert
	}
	
	// 找到镜像名称的结束位置
	end := start
	inImage := false
	for i, char := range line[start:] {
		if char == ' ' && inImage {
			end = start + i
			break
		}
		if char != ' ' {
			inImage = true
		}
	}
	
	if end == start {
		end = len(line)
	}
	
	return line[:end] + insert + line[end:]
}

// insertAfterNameParam 在--name参数后插入内容
func insertAfterNameParam(line, insert string) string {
	start := strings.Index(line, "--name")
	if start == -1 {
		return line + insert
	}
	
	// 找到--name参数值的结束位置
	start += len("--name")
	// 跳过空格
	for start < len(line) && line[start] == ' ' {
		start++
	}
	
	// 找到值的结束位置
	end := start
	for end < len(line) && line[end] != ' ' {
		end++
	}
	
	return line[:end] + insert + line[end:]
}

// restartOSDServices 重启OSD服务
func restartOSDServices(osds []OSDInfo, sshConfig *SSHConfig, fsid, method string, useSudo bool) {
	fmt.Println("\n重启OSD服务...")
	existing := detectExistingServices(osds, sshConfig, fsid, useSudo)
	cephadmAvail := commandAvailable(sshConfig, useSudo, "cephadm")
	cephAvail := false
	if cephadmAvail {
		cephAvail = cephCommandAvailableViaCephadm(sshConfig, useSudo)
	} else {
		cephAvail = commandAvailable(sshConfig, useSudo, "ceph")
	}
	if !cephadmAvail {
		fmt.Println("提示: 未检测到 cephadm，可用回退有限 (跳过 cephadm / ceph orch)。")
	}
	for _, osd := range osds {
		fmt.Printf("正在重启 OSD %d...\n", osd.ID)
		commands := buildRestartCommands(osd.ID, fsid, method, useSudo, existing[osd.ID], cephadmAvail, cephAvail)
		success := false
		triedSudoEscalation := false
		for i, c := range commands {
			// reset-failed + daemon-reload 预处理
			for _, svc := range existing[osd.ID] {
				if isServiceFailed(sshConfig, useSudo, svc) {
					_ = silentExec(sshConfig, useSudo, "systemctl reset-failed "+svc)
					_ = silentExec(sshConfig, useSudo, "systemctl daemon-reload")
				}
			}
			out, err := executeCommand(sshConfig, c)
			if err == nil {
				fmt.Printf("OSD %d 重启成功 (命令: %s)\n", osd.ID, c)
				success = true
				break
			}
			lower := strings.ToLower(string(out))
			if !useSudo && !triedSudoEscalation && (strings.Contains(lower, "permission") || strings.Contains(lower, "not permitted")) {
				sudoCmd := "sudo " + c
				fmt.Printf("检测到权限相关输出，尝试 sudo: %s\n", sudoCmd)
				out2, err2 := executeCommand(sshConfig, sudoCmd)
				triedSudoEscalation = true
				if err2 == nil {
					fmt.Printf("OSD %d 重启成功 (命令: %s)\n", osd.ID, sudoCmd)
					success = true
					break
				}
				fmt.Printf("sudo 回退失败 (OSD %d, 命令: %s) -> %v\n输出: %s\n", osd.ID, sudoCmd, err2, truncateOutput(string(out2), 300))
			}
			fmt.Printf("重启尝试失败 (OSD %d, 命令: %s) -> %v\n输出: %s\n", osd.ID, c, err, truncateOutput(string(out), 300))
			if i == len(commands)-1 {
				summarizeServiceState(osd.ID, sshConfig, fsid, useSudo, cephadmAvail)
			}
		}
		if !success {
			fmt.Printf("警告: OSD %d 所有重启方式失败。请检查上方状态摘要。\n", osd.ID)
		}
	}
	fmt.Println("所有OSD服务重启完成。")
}

// detectExistingServices 检测每个 OSD 可能的 systemd 服务名称中哪些存在
func detectExistingServices(osds []OSDInfo, sshConfig *SSHConfig, fsid string, useSudo bool) map[int][]string {
	result := make(map[int][]string)
	for _, osd := range osds {
		candidates := []string{
			fmt.Sprintf("ceph-%s@osd.%d.service", fsid, osd.ID),
			fmt.Sprintf("ceph-osd@%d.service", osd.ID),
		}
		for _, svc := range candidates {
			cmd := fmt.Sprintf("%ssystemctl status %s >/dev/null 2>&1", sudoPrefix(useSudo), svc)
			if _, err := executeCommand(sshConfig, cmd); err == nil {
				result[osd.ID] = append(result[osd.ID], svc)
			}
		}
	}
	return result
}

// buildRestartCommands 根据选择的方式构造重启命令列表（包含回退），只对存在的服务名生成 systemd 命令
func buildRestartCommands(id int, fsid, method string, sudo bool, existingServices []string, cephadmAvail, cephAvail bool) []string {
	pre := sudoPrefix(sudo)
	// 如果未检测到任何 systemd 服务，则后续仅使用非 systemd 回退（cephadm / orch）
	var cmds []string
	addSystemd := func() {
		for _, svc := range existingServices {
			cmds = append(cmds, pre+"systemctl restart "+svc, pre+"systemctl stop "+svc+" && "+pre+"systemctl start "+svc)
		}
	}
	switch method {
	case "systemd":
		addSystemd()
		// 如果没有 systemd 服务，添加提示性激活命令
		if len(existingServices) == 0 && cephadmAvail {
			cmds = append(cmds, pre+"cephadm shell -- ceph-volume lvm activate --all --force")
		}
	case "cephadm":
		if cephadmAvail {
			cmds = append(cmds, pre+"cephadm shell -- ceph-volume lvm activate --all --force")
		}
		addSystemd()
	case "orch":
		if cephadmAvail && cephAvail {
			cmds = append(cmds, pre+fmt.Sprintf("cephadm shell -- ceph orch restart osd.%d", id))
		}
		addSystemd()
	default:
		addSystemd()
		if cephadmAvail {
			cmds = append(cmds, pre+"cephadm shell -- ceph-volume lvm activate --all --force")
		}
		if cephadmAvail && cephAvail {
			cmds = append(cmds, pre+fmt.Sprintf("cephadm shell -- ceph orch restart osd.%d", id))
		}
	}
	// 去重
	seen := make(map[string]bool)
	var uniq []string
	for _, c := range cmds {
		if !seen[c] {
			uniq = append(uniq, c)
			seen[c] = true
		}
	}
	return uniq
}

// summarizeServiceState 输出服务状态摘要（存在的服务与最近日志行）
func summarizeServiceState(osdID int, sshConfig *SSHConfig, fsid string, useSudo bool, cephadmAvail bool) {
	fmt.Printf("[状态摘要 OSD %d]\n", osdID)
	listCmd := fmt.Sprintf("%ssystemctl list-units --type=service --all | grep -E 'ceph.*osd.%d' || true", sudoPrefix(useSudo), osdID)
	if out, err := executeCommand(sshConfig, listCmd); err == nil && len(out) > 0 {
		fmt.Printf("匹配的 systemd 服务:\n%s\n", truncateOutput(string(out), 500))
	}
	if cephadmAvail {
		orchCmd := fmt.Sprintf("%scephadm shell -- ceph orch ps --refresh | grep 'osd.%d' || true", sudoPrefix(useSudo), osdID)
		if out, err := executeCommand(sshConfig, orchCmd); err == nil && len(out) > 0 {
			fmt.Printf("orch 进程状态:\n%s\n", truncateOutput(string(out), 300))
		}
	}
	for _, svc := range []string{fmt.Sprintf("ceph-%s@osd.%d.service", fsid, osdID), fmt.Sprintf("ceph-osd@%d.service", osdID)} {
		if !serviceExists(sshConfig, useSudo, svc) { continue }
		logCmd := fmt.Sprintf("%sjournalctl -u %s -n 40 --no-pager 2>&1 | egrep -i 'error|failed|fatal|panic|cpuset|cgroup' || true", sudoPrefix(useSudo), svc)
		if out, err := executeCommand(sshConfig, logCmd); err == nil && len(out) > 0 {
			fmt.Printf("日志片段(%s):\n%s\n", svc, truncateOutput(string(out), 800))
		}
	}
}

func sudoPrefix(use bool) string { if use { return "sudo " }; return "" }

func serviceExists(sshConfig *SSHConfig, useSudo bool, svc string) bool {
	cmd := fmt.Sprintf("%ssystemctl status %s >/dev/null 2>&1", sudoPrefix(useSudo), svc)
	if _, err := executeCommand(sshConfig, cmd); err == nil { return true }
	return false
}

func isServiceFailed(sshConfig *SSHConfig, useSudo bool, svc string) bool {
	cmd := fmt.Sprintf("%ssystemctl is-failed %s 2>/dev/null || true", sudoPrefix(useSudo), svc)
	out, _ := executeCommand(sshConfig, cmd)
	return strings.Contains(strings.ToLower(string(out)), "failed")
}

func silentExec(sshConfig *SSHConfig, useSudo bool, cmd string) error {
	_, err := executeCommand(sshConfig, fmt.Sprintf("%s%s", sudoPrefix(useSudo), cmd))
	return err
}

func commandAvailable(sshConfig *SSHConfig, useSudo bool, cmd string) bool {
	check := fmt.Sprintf("%scommand -v %s >/dev/null 2>&1 || which %s >/dev/null 2>&1", sudoPrefix(useSudo), cmd, cmd)
	_, err := executeCommand(sshConfig, check)
	return err == nil
}

func cephCommandAvailableViaCephadm(sshConfig *SSHConfig, useSudo bool) bool {
	testCmd := fmt.Sprintf("%scephadm shell -- ceph -v >/dev/null 2>&1", sudoPrefix(useSudo))
	_, err := executeCommand(sshConfig, testCmd)
	return err == nil
}

// truncateOutput 截断过长输出
func truncateOutput(s string, max int) string {
	if len(s) <= max { return s }
	return s[:max] + "..." 
}

// parseOSDCPUMap 解析OSD到CPU的映射
func parseOSDCPUMap(osdCPUMap string, fsid string) (map[OSDInfo][]int, map[OSDInfo]int) {
	allocation := make(map[OSDInfo][]int)
	numaAssignment := make(map[OSDInfo]int)
	
	// 分割每个OSD的映射
	mappings := strings.Split(osdCPUMap, ",")
	for _, mapping := range mappings {
		// 分割OSD ID和CPU列表
		parts := strings.Split(mapping, ":")
		if len(parts) != 2 {
			continue
		}
		
		// 解析OSD ID
		osdID, err := strconv.Atoi(strings.TrimSpace(parts[0]))
		if err != nil {
			continue
		}
		
		// 解析CPU列表
		cpuList := strings.TrimSpace(parts[1])
		cpus := parseCPUListForOSD(cpuList)
		
		// 创建OSD信息
		osd := OSDInfo{ID: osdID, Path: fmt.Sprintf("/var/lib/ceph/%s/osd.%d", fsid, osdID)}
		
		// 存储分配信息
		allocation[osd] = cpus
		
		// 简单地将NUMA节点设置为0（实际应用中可能需要更复杂的逻辑）
		numaAssignment[osd] = 0
	}
	
	return allocation, numaAssignment
}

// parseCPUListForOSD 解析用于OSD的CPU列表
func parseCPUListForOSD(cpuList string) []int {
	var cpus []int
	parts := strings.Split(cpuList, ",")
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if strings.Contains(part, "-") {
			// 范围格式
			rangeParts := strings.Split(part, "-")
			if len(rangeParts) == 2 {
				start, _ := strconv.Atoi(rangeParts[0])
				end, _ := strconv.Atoi(rangeParts[1])
				for i := start; i <= end; i++ {
					cpus = append(cpus, i)
				}
			}
		} else {
			// 单个CPU
			if num, err := strconv.Atoi(part); err == nil {
				cpus = append(cpus, num)
			}
		}
	}
	return cpus
}
